package com.example.understandvideo.tool;

import com.google.gson.Gson;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;

import javax.swing.plaf.synth.SynthTextAreaUI;
import java.util.ArrayList;
import java.util.List;

public class ChatCompletion {
    public static List<String> imageVision(String content) {
        // 从环境变量中获取API密钥
        String apiKey = "b30fa44c-9e05-4ca9-a8a0-47f8fb8f81c8";

        // 创建ArkService实例
        ArkService arkService = ArkService.builder().apiKey(apiKey).build();

        // 初始化消息列表
        List<ChatMessage> chatMessages = new ArrayList<>();

        // 创建用户消息
        ChatMessage userMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER) // 设置消息角色为用户
                .content(content) // 设置消息内容
                .build();


        // 将用户消息添加到消息列表
        chatMessages.add(userMessage);

        // 创建聊天完成请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("doubao-1-5-vision-pro-32k-250115")// 需要替换为Model ID
                .messages(chatMessages) // 设置消息列表
                .build();

        System.out.println(chatCompletionRequest);
        // 发送聊天完成请求并打印响应
        List<String> result = null;
        try {
            Gson gson = new Gson();
            // 获取响应并打印每个选择的消息内容
            ChatCompletionResult chatCompletionResult = arkService.createChatCompletion(chatCompletionRequest);
            System.out.println(chatCompletionResult);
            chatCompletionResult
                    .getChoices()
                    .forEach(choice -> {
                        result.add(gson.toJson(choice.getMessage().getContent()));
                        System.out.println(choice.getMessage().getContent());
                    });
        } catch (Exception e) {
            System.out.println("请求失败: " + e.getMessage());
        } finally {
            // 关闭服务执行器
            arkService.shutdownExecutor();
        }
        return result;
    }
}
